/**
 * @file main.cpp
 * @brief AsyncVideoWriter类实现
 */

#include "../include/AsyncVideoWriter.h"

double now()
{
    timeval tv;
    gettimeofday(&tv, nullptr);
    return tv.tv_sec + tv.tv_usec * 1e-6;
}

template <typename TDataType>
WaitableQueue<TDataType>::WaitableQueue()
{
    node* new_node = new node();
    m_head.store(new_node);
    m_tail.store(new_node);
    m_length.store(0);
}

template <typename TDataType>
WaitableQueue<TDataType>::~WaitableQueue()
{
    set_stop();

    while (1)
    {
        node* old_head = m_head.load(std::memory_order_acquire);
        node* next = old_head->next.load(std::memory_order_acquire);

        if (old_head == m_head.load(std::memory_order_acquire))
        {
            if (next == nullptr)
            {
                delete old_head;
                break;
            }
            m_head.compare_exchange_strong(old_head, next,
                                           std::memory_order_release);
        }
        delete old_head;
    }
}

template <typename TDataType>
void WaitableQueue<TDataType>::push(const TDataType &value)
{
    node* new_node = new node(value);
    node* old_tail;

    while (1)
    {
        old_tail = m_tail.load(std::memory_order_acquire);
        node* next = old_tail->next.load(std::memory_order_acquire);

        if (old_tail != m_tail.load(std::memory_order_acquire))
            continue;

        if (next != nullptr)
        {
            m_tail.compare_exchange_weak(old_tail, next,
                                         std::memory_order_release);
            continue;
        }

        if (old_tail->next.compare_exchange_weak(next, new_node,
                                                 std::memory_order_release))
            break;
    }
    m_tail.compare_exchange_strong(old_tail, new_node,
                                   std::memory_order_release);
    m_length.fetch_add(1, std::memory_order_relaxed);
    m_condition_variable.notify_one();
}

template <typename TDataType>
void WaitableQueue<TDataType>::push(TDataType &&value)
{
    node* new_node = new node(std::move(value));
    node* old_tail;

    while (1)
    {
        old_tail = m_tail.load(std::memory_order_acquire);
        node* next = old_tail->next.load(std::memory_order_acquire);

        if (old_tail != m_tail.load(std::memory_order_acquire))
            continue;

        if (next != nullptr)
        {
            m_tail.compare_exchange_weak(old_tail, next,
                                         std::memory_order_release);
            continue;
        }

        if (old_tail->next.compare_exchange_weak(next, new_node,
                                                 std::memory_order_release))
            break;
    }
    m_tail.compare_exchange_strong(old_tail, new_node,
                                   std::memory_order_release);
    m_length.fetch_add(1, std::memory_order_relaxed);
    m_condition_variable.notify_one();
}

template <typename TDataType>
std::optional<TDataType> WaitableQueue<TDataType>::pop()
{
    while (1)
    {
        node* old_head = m_head.load(std::memory_order_acquire);
        node* old_tail = m_tail.load(std::memory_order_acquire);
        node* next = old_head->next.load(std::memory_order_acquire);

        if (old_head != m_head.load(std::memory_order_acquire))
            continue;

        if (old_head == old_tail)
        {
            if (next == nullptr)
                return std::nullopt;
            m_tail.compare_exchange_weak(old_tail, next,
                                         std::memory_order_release);
            continue;
        }

        std::shared_ptr<TDataType> result = next->data;

        if (m_head.compare_exchange_weak(old_head, next,
                                         std::memory_order_release))
        {
            m_length.fetch_sub(1, std::memory_order_relaxed);
            delete old_head;
            return std::move(*result);
        }
    }
}

template <typename TDataType>
std::optional<TDataType> WaitableQueue<TDataType>::wait_and_pop()
{
    std::unique_lock<std::mutex> lock(m_mutex);
    m_condition_variable.wait(lock, [this] { return !empty() || m_stop; });

    if (m_stop && empty())
        return std::nullopt;
    return pop();
}

template <typename TDataType>
void WaitableQueue<TDataType>::set_stop()
{
    m_stop = true;
    m_condition_variable.notify_all();
}

template <typename TDataType>
bool WaitableQueue<TDataType>::empty() const
{
    return m_length.load(std::memory_order_relaxed) == 0;
}

template <typename TDataType>
size_t WaitableQueue<TDataType>::size() const
{
    return m_length.load(std::memory_order_relaxed);
}

AsyncVideoWriter::AsyncVideoWriter()
    : m_image_size(0, 0),
      m_fps(30.0),
      m_fourcc(cv::VideoWriter::fourcc('M', 'J', 'P', 'G'))
{
}

AsyncVideoWriter::AsyncVideoWriter(
    const std::string &filename,
    int init_fourcc = cv::VideoWriter::fourcc('M', 'J', 'P', 'G'),
    double fps = 30.0, cv::Size frame_size = cv::Size(1280, 1024))
{
    this->m_file_name = filename;
    this->m_image_size = frame_size;
    this->m_fps = fps;
    this->m_fourcc = init_fourcc;
    this->m_is_running = true;
    this->m_should_stop = false;

    if (m_image_size.width == 0 || m_image_size.height == 0)
    {
        std::cerr << "in AsyncVideoWriter construction: init image size "
                     "should be no less than 0"
                  << std::endl;
        exit(-1);
    }

    // 写入视频线程
    m_writer_thread = std::thread(
        [this]()
        {
            m_video_writer.open(m_file_name, m_fourcc, this->m_fps,
                                m_image_size);

            if (!m_video_writer.isOpened())
            {
                std::cerr << "Failed to open VideoWriter" << std::endl;
                m_is_running = false;
                return;
            }

            // 当线程结束信号没来的时候一直循环
            while (!m_should_stop)
            {
                // wait_and_pop函数可以保证函数在没有拿到帧的时候会卡在这里而不是一直循环看队列是不是空的
                std::optional<cv::Mat> frame_opt = m_frame_queue.wait_and_pop();

                // 如果拿到std::nullopt(也就是没拿到帧)的逻辑
                if (!frame_opt.has_value())
                {
                    if (m_should_stop)
                        break;
                    continue;
                }

                // 使用移动语义获取frame
                cv::Mat frame = std::move(frame_opt.value());
                if (frame.empty())
                    break;

                m_video_writer << frame;

                // std::cout <<"writer << frame"<<std::endl;
            }

            // 线程结束信号结束之后，防止队列不为空，所以需要另一个循环
            while (1)
            {
                printf("left %lu frames\n", pendingFrames());
                std::optional<cv::Mat> frame_opt = m_frame_queue.pop();

                // 这里如果没有拿到图就说明队列为空，结束了
                if (!frame_opt.has_value())
                    break;

                cv::Mat frame = std::move(frame_opt.value());
                if (!frame.empty())
                    m_video_writer << frame;
            }

            m_video_writer.release();
            std::cout << "VideoWriter thread stop" << std::endl;

            // 将m_is_running置false，保证线程正常停止
            m_is_running = false;
        });
}

AsyncVideoWriter::~AsyncVideoWriter()
{
    if (m_is_running)
        release();

    if (m_writer_thread.joinable())
    {
        try
        {
            m_writer_thread.join();
        }
        catch (const std::exception &e)
        {
            std::cerr << "Failed to join writer thread in destructor"
                      << e.what() << '\n';
        }
    }
}

void AsyncVideoWriter::open(const std::string &filename,
                            int init_fourcc = cv::VideoWriter::fourcc('M', 'J',
                                                                      'P', 'G'),
                            double fps = 30.0,
                            cv::Size frame_size = cv::Size(1280, 1024))
{
    if (m_is_running)
        return;

    this->m_file_name = filename;
    this->m_image_size = frame_size;
    this->m_fps = fps;
    this->m_fourcc = init_fourcc;

    if (m_is_running || m_image_size.width == 0 || m_image_size.height == 0)
        return;

    m_is_running = true;
    m_should_stop = false;

    // 写入视频线程
    m_writer_thread = std::thread(
        [this]()
        {
            m_video_writer.open(m_file_name, m_fourcc, this->m_fps,
                                m_image_size);

            if (!m_video_writer.isOpened())
            {
                std::cerr << "Failed to open VideoWriter" << std::endl;
                m_is_running = false;
                return;
            }

            // 当线程结束信号没来的时候一直循环
            while (!m_should_stop)
            {
                // wait_and_pop函数可以保证函数在没有拿到帧的时候会卡在这里而不是一直循环看队列是不是空的
                std::optional<cv::Mat> frame_opt = m_frame_queue.wait_and_pop();

                // 如果拿到std::nullopt(也就是没拿到帧)的逻辑
                if (!frame_opt.has_value())
                {
                    if (m_should_stop)
                        break;
                    continue;
                }

                // 使用移动语义获取frame
                cv::Mat frame = std::move(frame_opt.value());
                if (frame.empty())
                    break;

                m_video_writer << frame;

                // std::cout <<"writer << frame"<<std::endl;
            }

            // 线程结束信号结束之后，防止队列不为空，所以需要另一个循环
            while (1)
            {
                printf("left %lu frames\n", pendingFrames());
                std::optional<cv::Mat> frame_opt = m_frame_queue.pop();

                // 这里如果没有拿到图就说明队列为空，结束了
                if (!frame_opt.has_value())
                    break;

                cv::Mat frame = std::move(frame_opt.value());
                if (!frame.empty())
                    m_video_writer << frame;
            }

            m_video_writer.release();
            std::cout << "VideoWriter thread stop" << std::endl;

            // 将m_is_running置false，保证线程正常停止
            m_is_running = false;
        });
    return;
}

void AsyncVideoWriter::write(const cv::Mat &frame)
{
    static double begin = now();
    if (!m_is_running || m_should_stop)
        return;

    if (now() - begin >= 1 / m_fps)
    {
        m_frame_queue.push(frame);
        begin = now();
    }
}

void AsyncVideoWriter::write(cv::Mat &&frame)
{
    static double begin = now();
    if (!m_is_running || m_should_stop)
        return;
    if (now() - begin >= 1 / m_fps)
    {
        m_frame_queue.push(std::move(frame));
        begin = now();
    }
}

void AsyncVideoWriter::operator<<(const cv::Mat &frame)
{
    static double begin = now();
    if (!m_is_running || m_should_stop)
        return;

    if (now() - begin >= 1 / m_fps)
    {
        m_frame_queue.push(frame);
        begin = now();
    }
}

void AsyncVideoWriter::operator<<(cv::Mat &&frame)
{
    static double begin = now();
    if (!m_is_running || m_should_stop)
        return;
    if (now() - begin >= 1 / m_fps)
    {
        m_frame_queue.push(std::move(frame));
        begin = now();
    }
}

void AsyncVideoWriter::release()
{
    if (!m_is_running)
        return;

    while (!m_frame_queue.empty())
        std::this_thread::sleep_for(std::chrono::milliseconds(10));

    m_should_stop = true;
    m_frame_queue.set_stop();

    m_frame_queue.push(cv::Mat());

    if (m_writer_thread.joinable())
    {
        try
        {
            m_writer_thread.join();
        }
        catch (const std::exception &e)
        {
            std::cerr << "Failed to join writer thread" << e.what() << '\n';
        }
    }

    m_is_running = false;
}

size_t AsyncVideoWriter::pendingFrames() const { return m_frame_queue.size(); }
